from PyPDF2 import PdfWriter
import PyPDF2

print('请选择您要的操作（1 合并  2 拆分） ：',  end='')
Go = input()
if Go == '1':
    print('请输入需要合并的pdf文档个数：', end='')
    numbers_pdf = int(input())
    merger = PdfWriter()
    list_pdf = []
    for number_pdf in range(numbers_pdf):
        print(f'请按顺序输入第 {number_pdf+1} 个pdf（例子：E:\BaiduNetdiskDownload\你好.pdf）：', end='')
        old_pdf = input()
        list_pdf.append(old_pdf)
    for pdf in list_pdf:
        merger.append(pdf)
    print('已合并完，请输入新文件名称（例子：aaa）不用添加扩展名：', end='')
    new_pdf_name = input()
    # 指定保存 .pdf 文件的完整路径和名称
    output_dir = "D:/"  # 替换为您的实际输出目录
    output_file_name = f"{new_pdf_name}_合并.pdf"
    output_path = output_dir + output_file_name
    merger.write(output_path)
    print('合并好了去查收吧')
    merger.close()

elif Go == '2':
        print('请输入您要查分的文件（例子：E:\BaiduNetdiskDownload\你好.pdf）：', end='')
        chaifen_pdf = input()
        open_pdf = open(chaifen_pdf, 'rb')
        read_pdf = PyPDF2.PdfReader(open_pdf)
        get_pdf_pages = len(read_pdf.pages)
        print(f"该文件共有{get_pdf_pages}页")
        print('---------------------------------------------------------')
        print('全部拆分请输入1 指定位置查分请输入2 ：', end='')
        split = input()

        if split == '1':
            for get_pdf_page in range(get_pdf_pages):
                page = read_pdf.pages[get_pdf_page]
                # 创建一个新的PdfWriter对象
                pdf_writer = PdfWriter()
                # 将当前页添加到新的PdfWriter对象中
                pdf_writer.add_page(page)
                # 获取用户输入的文件名，不包括路径
                pdf_file_name = chaifen_pdf.split('\\')[-1]  # 使用 \\
                pdf_base_name = pdf_file_name.split('.')[0]  # 去掉扩展名
                # 指定保存 .pdf 文件的完整路径和名称
                output_dir = "D:/"  # 替换为您的实际输出目录
                output_file_name = f"{pdf_base_name}_第{get_pdf_page + 1}页.pdf"
                output_path = output_dir + output_file_name
                pdf_writer.write(output_path)
                pdf_writer.close()
            print('拆分好了，快去看一下吧。')

        elif split == '2':
            print('请输入开始页 ：', end='')
            start_pages = int(input())
            print('请输入结束页 ：', end='')
            end_pages = int(input())
            # 创建一个新的PdfWriter对象
            pdf_writer = PdfWriter()
            for get_pdf_page in range(start_pages-1, end_pages):
                page = read_pdf.pages[get_pdf_page]
                # 将当前页添加到新的PdfWriter对象中
                pdf_writer.add_page(page)
            # 获取用户输入的文件名，不包括路径
            pdf_file_name = chaifen_pdf.split('\\')[-1]  # 使用 \\
            pdf_base_name = pdf_file_name.split('.')[0]  # 去掉扩展名
            # 指定保存 .pdf 文件的完整路径和名称
            output_dir = "D:/"  # 替换为您的实际输出目录
            output_file_name = f"{pdf_base_name}_第{start_pages}_{end_pages}页.pdf"
            output_path = output_dir + output_file_name
            pdf_writer.write(output_path)
            pdf_writer.close()
            print('拆分好了，快去看一下吧。')

